fix(bridge): derive position_ids from attention_mask for left-padded input - #1610
fix(bridge): derive position_ids from attention_mask for left-padded input#1610sohv wants to merge 2 commits into
Conversation
c37a8d8 to
cdc2af2
Compare
|
Thank you @jlarson4 for bringing these issues to my attention both were real. I reproduced these issues as you described and have pushed the revised code in commit cdc2af2. Instead of separately patching the two cases, I switched the derivation to Cached step. Reproduced: Interior gap. Reproduced: 6.752e+00 on gpt2 in compat mode against a 0.000e+00 unpadded control, with the gate correctly not firing. Rather than widen the predicate I dropped it — an all-ones mask gives One thing you did not flag that this surfaced: my pad-position convention was also wrong. I had Verification:
One correction to the PR description: the "this PR alone" row in the compounding table was 6.688155 and is now 5.396149, since pad positions inherit the previous index and that changes the still-unmasked aggregate #1607 produces. The "both" row is unchanged at 4.814578. |
…input TransformerBridge.forward() did not derive position_ids from a supplied attention_mask, so masked-out tokens silently shifted the absolute position of every real token after them — no error, no NaN, just wrong logits and a wrong loss. On gpt2 the loss for one prompt moved from 4.503170 unpadded to 11.154946 with three left pads, while HookedTransformer stays invariant (drift ~1e-06). transformer_bridge.py derived position_ids only for batched *list* input, so pre-tokenized tensors fell through to HF's plain arange and the offset was never removed. This reuses utils.get_offset_position_ids — the same helper PosEmbed and AbstractAttention already use — so the bridge shares HookedTransformer's position derivation rather than paralleling it. An explicitly supplied position_ids still wins. The derivation fires only when the mask actually moves an attended token off its default position, i.e. when some masked token precedes a real one. That covers left padding and interior mask gaps. Pure right padding and all-ones masks already agree with arange, so they are left alone: injecting position_ids there is a no-op at best, and breaks models whose forward does not accept the argument or which compute their own position streams (multimodal mRoPE). With a KV cache the mask spans past+new while input_ids holds only the new tokens, so the derived positions are sliced back to the tokens being passed. The bridge was also inconsistent with itself before this — the same batch gave different logits depending on whether it was passed as strings or token IDs (max |logit diff| 4.142e+01) — and enable_compatibility_mode(), which documents "HookedTransformer-equivalent numerics", diverged on left-padded input while matching exactly on unpadded input. Adds integration regression tests: logit invariance under both padding sides, the same property in compatibility mode, agreement with the shared helper, precedence of an explicit position_ids, interior mask gaps, a cached decode step, and that no position_ids are injected when the mask does not require it. Right-padding cases are controls that pass with and without the fix. They live in the integration tier because left padding produces a fully masked query row, which the Native attention path turns into NaN until the masked-softmax fix in TransformerLensOrg#1608 lands. Fixes TransformerLensOrg#1609. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
cdc2af2 to
32dd94d
Compare
|
CI caught a regression in my last push so this needed a third pass. Dropping the gate entirely was an over-correction. You asked me to widen it to any mask with a gap; I widened it to any mask at all, which made the derivation fire for cases that never needed it. Right padding and all-ones masks already agree with Fixed in 32dd94d: the derivation now fires only when the mask actually moves an attended token off its default position — some masked token precedes a real one. That is exactly left padding and interior gaps, and excludes right padding and all-ones masks.
One unrelated failure in that run: |
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Hey @sohv thanks for the detailed breakdown of your changes! I have a couple of additional follow ups based on this work.
Narrowing the predicate removed the tested symptoms (right padding, all-ones) without removing the mechanism: the bridge still injects a position_ids kwarg with no check that the target model accepts or wants it. Left padding and interior gaps still fire, and both of those reach architectures that break. A couple of line-specific comments to come for resolution:
| # The mask spans any cached prefix as well as the new tokens, so it is | ||
| # offset back to just the tokens actually being passed — matching how | ||
| # AbstractAttention/PosEmbed use past_kv_pos_offset. | ||
| if ( |
There was a problem hiding this comment.
For mRoPE architectures a left-padded batch now overrides the model's own 3-D derivation with a naive 2-D arange. HF only computes its rope index when position_ids is None (modeling_qwen2_5_vl.py:1249) and silently expands 2-D input across all three streams (lines 822-823) and for fixed-signature remote-code models like LLaDAModelLM the same injection raises TypeError, where base returned logits. Can injection be conditional on the target actually accepting and not already owning positions, the way output_attentions is guarded (bridge_core.py:1205-1218)?
| # there would be a no-op at best and unsupported by models whose | ||
| # forward does not take them at worst. | ||
| if bool(((_derived != _arange) & (attention_mask != 0)).any()): | ||
| kwargs["position_ids"] = _derived[ |
There was a problem hiding this comment.
The predicate is a whole-batch .any(), so a single left-padded row causes position_ids to be supplied for every row, including unpadded ones, which is what carries the mRoPE corruption into rows that had no padding at all. Is it possible to make the decision per-row?
…r row
The mask-derived position_ids from the previous commit were handed to every
model that hit the predicate. Two families break on that:
* Fixed-signature remote-code forwards (LLaDAModelLM) take neither
position_ids nor **kwargs, so a left-padded input raised TypeError where
the model used to return logits.
* Models that own their position derivation. mRoPE architectures build a
3-D index in get_rope_index, but only while position_ids is None; a
supplied 2-D tensor is silently broadcast across all three streams.
Left-padded multimodal input drifted 1.414e-04 off fresh HF.
* OPT's positional embedding consumes the mask and derives the same
positions, using its own convention for the padded slots, so injecting
replaced a correct derivation with a differing one.
_accepts_derived_position_ids() refuses all three, mirroring how
output_attentions is guarded in BridgeCore.run_with_cache. Their own
derivations already place positions on attended slots only, so deferring is
correct rather than a gap.
The predicate is now per row instead of a whole-batch .any(): rows that are
unpadded or purely right-padded keep arange verbatim, so one left-padded row
in a batch no longer perturbs its neighbours (8.223e-01 -> 0 on distilgpt2).
Also: cast the mask to long before deriving, since a float 0/1 mask produced
float positions and crashed the embedding lookup.
Cached decoding under left padding is fixed as a consequence -- prefilling
through the bridge and stepping now lands 1.755e-04 from the unpadded result,
against 2.844e+01 before. The previous test prefilled with raw HF and asserted
only shape and finiteness, so it modelled the one pattern that stays wrong.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2 participants
Description
Fixes #1609.
TransformerBridge.forward()did not deriveposition_idsfrom a suppliedattention_mask, so left-padded input silently got the wrong absolute positions — no error, no NaN, just wrong logits and a wrong loss.On gpt2, one prompt, mask supplied:
Right padding was never affected (drift ≤ 9.5e-07) — causality already protects it.
transformer_bridge.pyderivedposition_idsonly for batched list input, so pre-tokenized tensors fell through to HF's plainarangeand the padding offset was never removed. This extends the same correction that branch already applies. An explicitly suppliedposition_idsstill wins.Two consequences this also fixes:
4.142e+01).enable_compatibility_mode(), which documents "HookedTransformer-equivalent numerics", matched HT exactly on unpadded input (0.000e+00) but diverged on left-padded input.Tests
Adds
tests/integration/model_bridge/test_left_padding_positions.py: logit invariance under both padding sides, the same property in compatibility mode, and agreement between the derived and an explicitly suppliedposition_ids.Red-before / green-after: 9 passed with the fix, 5 failed without it. The right-padding cases are controls — they pass in both states, so the tests are specific to the bug rather than to padding in general.
These sit in the integration tier rather than the unit tier deliberately: left padding produces a fully masked query row, which the Native attention path turns into
NaNuntil the masked-softmax fix in #1608 lands, soboot_nativecannot express the property yet.Verification
tests/unit/model_bridge+tests/unit/test_tokenizer_padding_side.py: 3955 passed, 27 skipped, 10 xfailedpycln/isort/blackclean;mypycleanRelationship to #1607 / #1608
Independent bugs on the same path that compound. This is measured across four states (gpt2, aggregate loss on a left-padded batch; HT reference
4.814578):dev-4.xThis PR fixes the logits; #1608 fixes the loss aggregation. Batched loss is only correct with both, so reviewing this one in isolation will still show a wrong aggregate. No file overlap, so they merge in either order.
Type of change
Checklist: